package com.sinnatec.crypto.sm2;

import org.bouncycastle.asn1.gm.GMObjectIdentifiers;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.jce.spec.ECParameterSpec;

import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class SM2Util {

    static {
        if (null == Security.getProvider(BouncyCastleProvider.PROVIDER_NAME)) {
            Security.addProvider(new BouncyCastleProvider());
        }
    }

    private SM2Util(){

    }

    /**
     * 根据publicKey对原始数据data，使用SM2加密
     *
     * @param content
     * @param publicKeyStr
     * @return
     */
    public static String encrypt(String content, String publicKeyStr) {
        ECPublicKeyParameters localECPublicKeyParameters = null;
        PublicKey publicKey = createPublicKey(publicKeyStr);
        if (publicKey instanceof BCECPublicKey) {
            BCECPublicKey localECPublicKey = (BCECPublicKey) publicKey;
            ECParameterSpec localECParameterSpec = localECPublicKey.getParameters();
            ECDomainParameters localECDomainParameters = new ECDomainParameters(localECParameterSpec.getCurve(),
                    localECParameterSpec.getG(), localECParameterSpec.getN());
            localECPublicKeyParameters = new ECPublicKeyParameters(localECPublicKey.getQ(), localECDomainParameters);
        }
        SM2Engine localSM2Engine = new SM2Engine();
        localSM2Engine.init(true, new ParametersWithRandom(localECPublicKeyParameters, new SecureRandom()));

        try {
            byte [] data = content.getBytes(StandardCharsets.UTF_8);
            byte[] encryptdBytes = localSM2Engine.processBlock(data, 0, data.length);
            return Base64.getEncoder().encodeToString(encryptdBytes);
        } catch (InvalidCipherTextException e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 根据privateKey对加密数据encodedata，使用SM2解密
     *
     * @param data
     * @param privateKeyStr
     * @return
     */
    public static String decrypt(String data,  String privateKeyStr) {
        SM2Engine localSM2Engine = new SM2Engine();
        PrivateKey privateKey = createPrivateKey(privateKeyStr);
        BCECPrivateKey sm2PriK = (BCECPrivateKey) privateKey;
        ECParameterSpec localECParameterSpec = sm2PriK.getParameters();
        ECDomainParameters localECDomainParameters = new ECDomainParameters(localECParameterSpec.getCurve(),
                localECParameterSpec.getG(), localECParameterSpec.getN());
        ECPrivateKeyParameters localECPrivateKeyParameters = new ECPrivateKeyParameters(sm2PriK.getD(),
                localECDomainParameters);
        localSM2Engine.init(false, localECPrivateKeyParameters);
        try {
            byte [] byteData = Base64.getDecoder().decode(data);
            byte [] decryptedData = localSM2Engine.processBlock(Base64.getDecoder().decode(data), 0, byteData.length);
            return new String(decryptedData,StandardCharsets.UTF_8);
        } catch (InvalidCipherTextException e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 私钥签名
     *
     * @param content
     * @param privateKeyStr
     * @return
     * @throws Exception
     */
    public static String signByPrivateKey(String content, String privateKeyStr) throws NoSuchAlgorithmException, NoSuchProviderException, InvalidKeyException, SignatureException {
        PrivateKey privateKey = createPrivateKey(privateKeyStr);
        byte[] data = content.getBytes(StandardCharsets.UTF_8);
        Signature sig = Signature.getInstance(GMObjectIdentifiers.sm2sign_with_sm3.toString(), BouncyCastleProvider.PROVIDER_NAME);
        sig.initSign(privateKey);
        sig.update(data);
        return Base64.getEncoder().encodeToString(sig.sign());
    }

    /**
     * 公钥验签
     *
     * @param content
     * @param publicKeyStr
     * @param signature
     * @return
     * @throws Exception
     */
    public static boolean verifyByPublicKey(String content, String publicKeyStr, String signature) throws NoSuchAlgorithmException, NoSuchProviderException, InvalidKeyException, SignatureException {
        Signature sig = Signature.getInstance(GMObjectIdentifiers.sm2sign_with_sm3.toString(), BouncyCastleProvider.PROVIDER_NAME);
        PublicKey publicKey = createPublicKey(publicKeyStr);
        sig.initVerify(publicKey);
        sig.update(content.getBytes(StandardCharsets.UTF_8));
        return sig.verify(Base64.getDecoder().decode(signature));
    }

    /**
     * 生成国密公私钥对
     *
     * @return
     * @throws Exception
     */
    public static SM2KeyPair generateKey() throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
        KeyPairGenerator keyPairGenerator = null;
        ECGenParameterSpec sm2Spec = new ECGenParameterSpec("sm2p256v1");
        keyPairGenerator = KeyPairGenerator.getInstance("EC", Security.getProvider(BouncyCastleProvider.PROVIDER_NAME));
        keyPairGenerator.initialize(sm2Spec);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        PrivateKey privateKey = keyPair.getPrivate();
        PublicKey publicKey = keyPair.getPublic();
        return new SM2KeyPair(Base64.getEncoder().encodeToString(privateKey.getEncoded()),Base64.getEncoder().encodeToString(publicKey.getEncoded()));
    }

    /**
     * 将Base64转码的公钥串，转化为公钥对象
     *
     * @param publicKey
     * @return
     */
    private static PublicKey createPublicKey(String publicKey) {
        PublicKey publickey = null;
        try {
            X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(Base64.getDecoder().decode(publicKey));
            KeyFactory keyFactory = KeyFactory.getInstance("EC", new BouncyCastleProvider());
            publickey = keyFactory.generatePublic(publicKeySpec);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return publickey;
    }

    /**
     * 将Base64转码的私钥串，转化为私钥对象
     *
     * @param privateKey
     * @return
     */
    private static PrivateKey createPrivateKey(String privateKey) {
        PrivateKey publickey = null;
        try {
            PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(Base64.getDecoder().decode(privateKey));
            KeyFactory keyFactory = KeyFactory.getInstance("EC", new BouncyCastleProvider());
            publickey = keyFactory.generatePrivate(pkcs8EncodedKeySpec);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return publickey;
    }

}
